from ast import main
import os
import numpy as np
import matplotlib.pyplot as plt
from regex import D


def create_folder_if_not_exists(folder_path):
    if not os.path.exists(folder_path):
        try:
            # 创建文件夹
            os.makedirs(folder_path)
            print(f"Folder '{folder_path}' created successfully.")
        except OSError as e:
            print(f"Error: Unable to create folder '{folder_path}'. {e}")
            return False
    return True


def problemA():
    # 读取文本文件
    data0 = np.loadtxt("txtfile/Af0.txt")
    data1 = np.loadtxt("txtfile/Af1.txt")
    data2 = np.loadtxt("txtfile/Af2.txt")
    data3 = np.loadtxt("txtfile/Af3.txt")
    data4 = np.loadtxt("txtfile/Af4.txt")
    data5 = np.loadtxt("txtfile/Af5.txt")

    x = data0[:, 0]
    y0 = data0[:, 1]
    y1 = data1[:, 1]
    y2 = data2[:, 1]
    y3 = data3[:, 1]
    y4 = data4[:, 1]
    y5 = data5[:, 1]

    # 绘制两个函数的图像
    plt.figure(1, figsize=(10, 8))
    plt.plot(x, y0, label="Function exact")
    plt.plot(x, y1, label="Function n=6")
    plt.plot(x, y2, label="Function n=11")
    plt.plot(x, y3, label="Function n=21")
    plt.plot(x, y4, label="Function n=41")
    plt.plot(x, y5, label="Function n=81")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.title("problem A")
    plt.legend()
    plt.grid(True)
    plt.savefig("picfile/problemA.png")
    plt.show()


def problemC():
    # 读取文本文件
    data0 = np.loadtxt("txtfile/Cf0.txt")
    data1 = np.loadtxt("txtfile/Cf1.txt")
    data2 = np.loadtxt("txtfile/Cf2.txt")

    x = data0[:, 0]
    y0 = data0[:, 1]
    y1 = data1[:, 1]
    y2 = data2[:, 1]

    # 绘制两个函数的图像
    plt.figure(1, figsize=(10, 8))
    plt.plot(x, y0, label="Function exact")
    plt.plot(x, y1, label="Function thm3.57")
    plt.plot(x, y2, label="Function thm3.58")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.title("problem C")
    plt.legend()
    plt.grid(True)
    plt.savefig("picfile/problemC.png")
    plt.show()


def problemE():
    # 读取文本文件

    data0 = np.loadtxt("txtfile/Ef0.txt")
    data1 = np.loadtxt("txtfile/Ef1.txt")
    data2 = np.loadtxt("txtfile/Ef2.txt")
    data3 = np.loadtxt("txtfile/Ef3.txt")

    x0 = data0[:, 0]
    y0 = data0[:, 1]
    x1 = data1[:, 0]
    y1 = data1[:, 1]
    x2 = data2[:, 0]
    y2 = data2[:, 1]
    x3 = data3[:, 0]
    y3 = data3[:, 1]

    # 绘制两个函数的图像
    plt.figure(1, figsize=(10, 8))
    plt.plot(x0, y0, label="Function exact")
    plt.plot(x1, y1, label="Function n=10")
    plt.plot(x2, y2, label="Function n=40")
    plt.plot(x3, y3, label="Function n=160")
    plt.xlabel("x")
    plt.ylabel("y")
    plt.title("problem E")
    plt.legend()
    plt.grid(True)
    plt.savefig("picfile/problemE.png")
    plt.show()


if __name__ == '__main__':
    create_folder_if_not_exists("picfile")
    problemA()
    problemC()
    problemE()
